// OpenCV
#include <opencv2/core.hpp>
#include <opencv2/core/base.hpp>
#include <opencv2/core/mat.hpp>
#include <opencv2/core/types.hpp>
#include <opencv2/imgproc.hpp>

// STD
#include <algorithm>  // <algorithm>提供各种算法，例如查找、排序等常用算法函数，在这里可能用于对数据结构（如容器中的元素）进行操作。
#include <cmath>  // <cmath>用于数学计算相关的函数，比如三角函数、幂运算等，在涉及角度计算、数值比例计算等地方可能会用到。
#include <vector> // <vector>用于定义和操作动态大小的数组，即向量容器，代码中很多地方用它来存储图像中的轮廓、检测到的灯条、装甲等对象集合。

#include "armor_detector/detector.hpp"
#include "auto_aim_interfaces/msg/debug_armor.hpp"
#include "auto_aim_interfaces/msg/debug_light.hpp"

namespace rm_auto_aim
{
  /* 定义Detector类的构造函数，它接受几个参数：
   *  bin_thres：可能是用于图像二值化处理的阈值参数，用于将图像转换为二值图像（黑白图像）的阈值设定。
   *  color：可能表示要检测的目标颜色相关的标识，用于后续筛选符合特定颜色的目标。
   *  l：应该是LightParams类型的参数，可能包含了与灯条检测相关的各种参数设定，例如灯条的形状、角度、填充比例等方面的阈值等参数。
   *  a：应该是ArmorParams类型的参数，可能包含了与装甲识别相关的各种参数设定，比如装甲组成的灯条之间的距离、角度、长度比例等方面的阈值参数。
   *  通过初始化列表的方式初始化类中的成员变量，将传入的参数值赋给对应的成员变量。
  */
  Detector::Detector(const int & bin_thres, const int & color, const LightParams & l, const ArmorParams & a)
  : binary_thres(bin_thres), detect_color(color), l(l), a(a)
  {
    
  }

  /* Detector类的成员函数detect，用于执行整个目标检测流程，它接受一个cv::Mat类型的输入图像参数input，代表要进行目标检测的原始图像：
   *  首先调用preprocessImage函数对输入图像进行预处理，将处理后的二值图像结果存储在成员变量binary_img中。
   *  接着调用findLights函数，基于原始图像和二值图像来查找图像中的灯条信息，将找到的灯条信息存储在成员变量lights_中，lights_ 应该是一个存储Light类型对象的向量容器。
   *  然后调用matchLights函数，根据找到的灯条信息尝试匹配出可能的装甲信息，将匹配到的装甲信息存储在成员变量armors_中，armors_ 同样是一个存储Armor类型对象的向量容器。
   *  如果检测到的装甲信息不为空，即找到了至少一个装甲，那么调用classifier（应该是一个用于分类识别的对象）的extractNumbers和classify函数，对装甲上的数字进行提取和分类识别操作，具体功能依赖于classifier对象的实现。
   *  最后返回检测到的装甲信息，即armors_向量容器，外部调用者可以获取到检测到的装甲相关结果。
  */
  std::vector<Armor> Detector::detect(const cv::Mat & input)
  {
    binary_img = preprocessImage(input);
    lights_ = findLights(input, binary_img);
    armors_ = matchLights(lights_);

    if (!armors_.empty()) 
    {
      classifier->extractNumbers(input, armors_);
      classifier->classify(armors_);
    }

    return armors_;
  }

  /* Detector类的成员函数preprocessImage，用于对输入的彩色图像（RGB格式）进行预处理：
   *  首先创建一个cv::Mat类型的变量gray_img，用于存储转换后的灰度图像。
   *  调用cv::cvtColor函数，将输入的RGB格式图像rgb_img转换为灰度图像，转换的方式按照cv::COLOR_RGB2GRAY指定的颜色空间转换模式进行，结果存储在gray_img中。
   *  接着创建另一个cv::Mat类型的变量binary_img，用于存储二值化后的图像。
   *  调用cv::threshold函数，根据成员变量binary_thres设定的阈值，对灰度图像gray_img进行二值化处理，将像素值大于阈值的像素设置为255（白色），小于等于阈值的像素设置为0（黑色），二值化的模式采用cv::THRESH_BINARY，处理后的结果存储在binary_img中。
   *  最后返回二值化后的图像binary_img，供后续的检测流程使用。
  */
  cv::Mat Detector::preprocessImage(const cv::Mat & rgb_img)
  {
    cv::Mat gray_img;
    cv::cvtColor(rgb_img, gray_img, cv::COLOR_RGB2GRAY);

    cv::Mat binary_img;
    cv::threshold(gray_img, binary_img, binary_thres, 255, cv::THRESH_BINARY);

    return binary_img;
  }

  /* Detector类的成员函数findLights，用于在给定的彩色图像和对应的二值图像中查找灯条信息：
     首先定义两个向量容器contours和hierarchy，用于存储通过cv::findContours函数查找得到的图像轮廓信息，contours存储轮廓的点集，hierarchy存储轮廓的层次结构信息，调用cv::findContours函数按照cv::RETR_EXTERNAL（只检测最外层轮廓）和cv::CHAIN_APPROX_SIMPLE（压缩水平、垂直和斜向的线段，只保留端点信息）的模式来查找二值图像binary_img中的轮廓信息。
     创建一个存储Light类型对象的向量容器lights，用于后续存储找到的灯条信息，同时清空之前存储在debug_lights.data中的调试数据（debug_lights应该是用于记录灯条相关调试信息的对象）。
     遍历查找到的每个轮廓contour：
     如果轮廓的点数小于5，认为不符合灯条的基本形状要求，直接跳过该轮廓继续下一个轮廓的处理。
     对于符合条件的轮廓，计算其外接矩形b_rect（使用cv::boundingRect函数）和最小外接矩形r_rect（使用cv::minAreaRect函数），用于描述轮廓的大致形状和位置信息。
     创建一个与外接矩形b_rect大小相同的全零掩膜图像mask（类型为CV_8UC1，即单通道8位无符号整数类型的图像），并根据轮廓的点集信息构建一个相对外接矩形坐标原点的点集mask_contour，用于后续在掩膜图像上填充对应轮廓区域。
     使用cv::fillPoly函数将mask_contour所表示的轮廓区域在掩膜图像mask上填充为白色（像素值为255），然后通过cv::findNonZero函数找到掩膜图像中不为零（白色）的像素点坐标，存储在points向量中。
     根据points向量中点的数量与最小外接矩形的面积比例，判断该轮廓区域是否满足最小填充比例要求（通过与LightParams类型的成员变量l中的min_fill_ratio进行比较），将结果存储在is_fill_rotated_rect变量中，用于后续筛选灯条。
     使用cv::fitLine函数对points中的点集拟合一条直线，得到直线的参数return_param，根据直线参数尝试计算灯条的上下端点坐标top和bottom以及倾斜角度angle_k（通过不同的条件判断和数学计算来确定端点坐标和角度，当直线为水平或垂直等特殊情况有特定的计算方式，非特殊情况则通过直线斜率等计算来确定端点和角度），并以此构建一个Light类型的对象light，它包含了灯条的位置、形状、角度等相关信息。
     调用isLight函数判断构建的light对象是否符合灯条的其他条件要求（比如长宽比、倾斜角度等条件），并且判断是否满足最小填充比例要求is_fill_rotated_rect，如果都满足，则进一步进行一些处理：
     先将light对象赋值给rect变量（可能是为了后续方便进行边界检查等操作），然后检查该灯条区域是否在原始彩色图像rgb_img的范围内（通过坐标范围判断，避免越界访问导致断言失败等问题），如果在范围内，则对灯条区域内的像素进行颜色分析：
     定义sum_r和sum_b变量，分别用于统计灯条区域内红色像素和蓝色像素的总和，通过获取灯条对应的感兴趣区域roi（使用rgb_img(rect)方式获取图像的子区域），然后遍历roi区域内的每个像素点，使用cv::pointPolygonTest函数判断像素点是否在轮廓内部，如果在内部，则根据像素点的颜色通道值（对于RGB图像，[0]表示红色通道，[2]表示蓝色通道）累加到对应的sum_r或sum_b变量中。
     根据红色像素和蓝色像素总和的比较结果，确定灯条的颜色（如果红色像素总和大于蓝色像素总和，认为是红色灯条，否则为蓝色灯条），并将该灯条对象添加到lights向量容器中，最终将存储了所有检测到的灯条信息的lights向量容器返回。
  */
  std::vector<Light> Detector::findLights(const cv::Mat & rgb_img, const cv::Mat & binary_img)
  {
    using std::vector;
    vector<vector<cv::Point>> contours;
    vector<cv::Vec4i> hierarchy;
    cv::findContours(binary_img, contours, hierarchy, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE);

    vector<Light> lights;
    this->debug_lights.data.clear();

    for (const auto & contour : contours) 
    {
      if (contour.size() < 5) continue;

      auto b_rect = cv::boundingRect(contour);
      auto r_rect = cv::minAreaRect(contour);
      cv::Mat mask = cv::Mat::zeros(b_rect.size(), CV_8UC1);
      std::vector<cv::Point> mask_contour;
      for (const auto & p : contour) 
      {
        mask_contour.emplace_back(p - cv::Point(b_rect.x, b_rect.y));
      }

      cv::fillPoly(mask, {mask_contour}, 255);
      std::vector<cv::Point> points;
      cv::findNonZero(mask, points);
      // points / rotated rect area
      bool is_fill_rotated_rect = points.size() / (r_rect.size.width * r_rect.size.height) > l.min_fill_ratio;
      cv::Vec4f return_param;
      cv::fitLine(points, return_param, cv::DIST_L2, 0, 0.01, 0.01);
      cv::Point2f top, bottom;
      double angle_k;
      if (int(return_param[0] * 100) == 100 || int(return_param[1] * 100) == 0) 
      {
        top = cv::Point2f(b_rect.x + b_rect.width / 2, b_rect.y);
        bottom = cv::Point2f(b_rect.x + b_rect.width / 2, b_rect.y + b_rect.height);
        angle_k = 0;
      } 
      else 
      {
        auto k = return_param[1] / return_param[0];
        auto b = (return_param[3] + b_rect.y) - k * (return_param[2] + b_rect.x);
        top = cv::Point2f((b_rect.y - b) / k, b_rect.y);
        bottom = cv::Point2f((b_rect.y + b_rect.height - b) / k, b_rect.y + b_rect.height);
        angle_k = std::atan(k) / CV_PI * 180 - 90;
        if (angle_k > 90) 
        {
          angle_k = 180 - angle_k;
        }
      }
      auto light = Light(b_rect, top, bottom, points.size(), angle_k);

      if (isLight(light) && is_fill_rotated_rect) 
      {
        auto rect = light;
        if (  // Avoid assertion failed
          0 <= rect.x && 0 <= rect.width && rect.x + rect.width <= rgb_img.cols && 0 <= rect.y &&
          0 <= rect.height && rect.y + rect.height <= rgb_img.rows) 
        {
          int sum_r = 0, sum_b = 0;
          auto roi = rgb_img(rect);
          // Iterate through the ROI
          for (int i = 0; i < roi.rows; i++) 
          {
            for (int j = 0; j < roi.cols; j++) 
            {
              if (cv::pointPolygonTest(contour, cv::Point2f(j + rect.x, i + rect.y), false) >= 0) 
              {
                // if point is inside contour
                sum_r += roi.at<cv::Vec3b>(i, j)[0];
                sum_b += roi.at<cv::Vec3b>(i, j)[2];
              }
            }
          }
          // Sum of red pixels > sum of blue pixels ?
          light.color = sum_r > sum_b ? RED : BLUE;
          lights.emplace_back(light);
        }
      }
    }

    return lights;
  }

  /* Detector类的成员函数isLight，用于判断给定的Light类型对象是否符合灯条的条件要求：
     首先计算灯条的宽长比（短边长度除以长边长度），将结果存储在ratio变量中，然后通过与LightParams类型的成员变量l中的min_ratio和max_ratio进行比较，判断宽长比是否在合理范围内，将结果存储在ratio_ok变量中，用于衡量灯条形状是否符合要求。
     判断灯条的倾斜角度light.tilt_angle是否小于LightParams类型成员变量l中的max_angle，将结果存储在angle_ok变量中，用于衡量灯条角度是否符合要求。
     只有当宽长比和倾斜角度都符合要求（即ratio_ok和angle_ok都为true）时，才认为该对象是一个有效的灯条，将结果存储在is_light变量中。
     创建一个auto_aim_interfaces::msg::DebugLight类型的消息结构体对象light_data，用于记录灯条的调试信息，将灯条的中心横坐标、宽长比、倾斜角度以及是否为有效灯条的判断结果等信息填充到light_data对象中，然后将该调试信息对象添加到debug_lights.data向量容器中（用于后续调试输出等操作）。
     最后返回is_light变量的值，即判断该Light对象是否为有效灯条的结果，外部调用者可以根据返回值进行后续处理。
  */  
  bool Detector::isLight(const Light & light)
  {
    // The ratio of light (short side / long side)
    float ratio = light.width / light.length;
    bool ratio_ok = l.min_ratio < ratio && ratio < l.max_ratio;

    bool angle_ok = light.tilt_angle < l.max_angle;

    bool is_light = ratio_ok && angle_ok;

    // Fill in debug information
    auto_aim_interfaces::msg::DebugLight light_data;
    light_data.center_x = light.center.x;
    light_data.ratio = ratio;
    light_data.angle = light.tilt_angle;
    light_data.is_light = is_light;
    this->debug_lights.data.emplace_back(light_data);

    return is_light;
  }

  /* 遍历输入的所有灯条（lights）的两两组合，通过颜色筛选、外接矩形内是否包含其他灯条判断以及是否构成有效装甲的判断等条件筛选，找出所有符合条件的装甲信息，最后返回包含这些装甲信息的向量容器。 */
  std::vector<Armor> Detector::matchLights(const std::vector<Light> & lights)
  {
    std::vector<Armor> armors;
    this->debug_armors.data.clear();

    // Loop all the pairing of lights
    for (auto light_1 = lights.begin(); light_1 != lights.end(); light_1++) 
    {
      for (auto light_2 = light_1 + 1; light_2 != lights.end(); light_2++) 
      {
        // 首先进行颜色筛选，如果两个灯条（light_1和light_2）的颜色与构造函数传入的要检测的目标颜色detect_color不一致，就跳过本次循环，不考虑这一对灯条的组合，因为不符合期望的颜色要求。
        if (light_1->color != detect_color || light_2->color != detect_color) continue;
        // 调用containLight函数，判断由这两个灯条（light_1和light_2）构成的外接矩形区域内是否包含其他灯条，如果包含其他灯条，则不符合装甲组成的条件，跳过本次循环，不进行后续关于这对灯条是否构成装甲的判断。
        if (containLight(*light_1, *light_2, lights)) 
        {
          continue;
        }
        /*  调用isArmor函数，判断这两个灯条是否能构成有效的装甲，
            如果返回的类型不是ArmorType::INVALID（说明是有效的装甲类型，比如LARGE或者SMALL类型），
            则创建一个Armor类型的对象armor，将其类型设置为判断得到的有效类型type，
            然后把这个armor对象添加到armors向量容器中，armors容器用于存储所有检测到的、符合条件的装甲信息。
        */
        auto type = isArmor(*light_1, *light_2);
        if (type != ArmorType::INVALID) 
        {
          auto armor = Armor(*light_1, *light_2);
          armor.type = type;
          armors.emplace_back(armor);
        }
      }
    }

    return armors;
  }

  // Check if there is another light in the boundingRect formed by the 2 lights
  bool Detector::containLight(
    const Light & light_1, const Light & light_2, const std::vector<Light> & lights)
  {
    /*  首先根据传入的两个灯条（light_1和light_2）的上下端点坐标信息，构建一个包含这四个点的向量points，
        然后使用cv::boundingRect函数计算出这四个点构成的最小外接矩形bounding_rect，这个外接矩形就是由这两个灯条所界定的大致区域。
    */
    auto points = std::vector<cv::Point2f>{light_1.top, light_1.bottom, light_2.top, light_2.bottom};
    auto bounding_rect = cv::boundingRect(points);

    for (const auto & test_light : lights) 
    {
      /* 遍历所有的灯条（lights向量中的每个元素），对于每个灯条test_light，
         如果它的中心坐标与light_1或light_2的中心坐标相同，说明是同一个灯条，直接跳过本次循环，不进行后续包含性判断，
         因为我们要判断的是除这两个灯条之外的其他灯条是否在该外接矩形内。
      */
      if (test_light.center == light_1.center || test_light.center == light_2.center) continue;
      /*  检查当前遍历到的灯条test_light的上顶点、下顶点或者中心坐标是否在外接矩形bounding_rect内部，
          如果有任何一个点在内部，就说明在由light_1和light_2构成的外接矩形区域内包含了其他灯条，此时返回true，表示存在这种包含情况。
      */
      if (bounding_rect.contains(test_light.top) || bounding_rect.contains(test_light.bottom) ||
          bounding_rect.contains(test_light.center)) 
      {
        return true;
      }
    }

    return false;
  }

  ArmorType Detector::isArmor(const Light & light_1, const Light & light_2)
  {
    // Ratio of the length of 2 lights (short side / long side)
    float light_length_ratio = light_1.length < light_2.length ? light_1.length / light_2.length
                                                              : light_2.length / light_1.length;
    bool light_ratio_ok = light_length_ratio > a.min_light_ratio;
    /*  计算两个灯条长度的比例（取短边长度除以长边长度），
        然后与ArmorParams类型的成员变量a中的min_light_ratio进行比较，
        判断该长度比例是否大于设定的最小灯条长度比例阈值，
        如果大于则说明长度比例方面符合构成装甲的条件
    */
    // Distance between the center of 2 lights (unit : light length)
    float avg_light_length = (light_1.length + light_2.length) / 2;
    float center_distance = cv::norm(light_1.center - light_2.center) / avg_light_length;
    bool center_distance_ok = (a.min_small_center_distance <= center_distance &&
                              center_distance < a.max_small_center_distance) ||
                              (a.min_large_center_distance <= center_distance &&
                              center_distance < a.max_large_center_distance);

    // Angle of light center connection
    cv::Point2f diff = light_1.center - light_2.center;
    float angle = std::abs(std::atan(diff.y / diff.x)) / CV_PI * 180;
    bool angle_ok = angle < a.max_angle;

    bool is_armor = light_ratio_ok && center_distance_ok && angle_ok;

    // Judge armor type
    ArmorType type;
    if (is_armor) 
    {
      type = center_distance > a.min_large_center_distance ? ArmorType::LARGE : ArmorType::SMALL;
    } 
    else 
    {
      type = ArmorType::INVALID;
    }

    // Fill in debug information
    auto_aim_interfaces::msg::DebugArmor armor_data;
    armor_data.type = ARMOR_TYPE_STR[static_cast<int>(type)];
    armor_data.center_x = (light_1.center.x + light_2.center.x) / 2;
    armor_data.light_ratio = light_length_ratio;
    armor_data.center_distance = center_distance;
    armor_data.angle = angle;
    this->debug_armors.data.emplace_back(armor_data);

    return type;
  }

  // 获取所有检测到的装甲对应的数字图像，并将它们按垂直方向拼接成一个整体图像返回，如果没有检测到装甲则返回一个空白图像。
  cv::Mat Detector::getAllNumbersImage()
  {
    if (armors_.empty()) 
    {
      return cv::Mat(cv::Size(20, 28), CV_8UC1);
    } 
    else 
    {
      std::vector<cv::Mat> number_imgs;
      number_imgs.reserve(armors_.size());
      for (auto & armor : armors_) 
      {
        number_imgs.emplace_back(armor.number_img);
      }
      cv::Mat all_num_img;
      cv::vconcat(number_imgs, all_num_img);
      return all_num_img;
    }
  }

  // 在给定的图像 img 上绘制出检测到的灯条、装甲的相关图形以及显示装甲的分类结果等信息，实现对目标检测结果的可视化展示，方便直观地查看检测效果和相关信息。
  void Detector::drawResults(cv::Mat & img)
  {
    // Draw Lights
    for (const auto & light : lights_) 
    {
      cv::circle(img, light.top, 3, cv::Scalar(255, 255, 255), 1);
      cv::circle(img, light.bottom, 3, cv::Scalar(255, 255, 255), 1);
      auto line_color = light.color == RED ? cv::Scalar(255, 255, 0) : cv::Scalar(255, 0, 255);
      cv::line(img, light.top, light.bottom, line_color, 1);
    }

    // Draw armors
    for (const auto & armor : armors_) 
    {
      cv::line(img, armor.left_light.top, armor.right_light.bottom, cv::Scalar(0, 255, 0), 2);
      cv::line(img, armor.left_light.bottom, armor.right_light.top, cv::Scalar(0, 255, 0), 2);
    }

    // Show numbers and confidence
    for (const auto & armor : armors_) 
    {
      cv::putText(img, armor.classfication_result, armor.left_light.top, cv::FONT_HERSHEY_SIMPLEX, 0.8, cv::Scalar(0, 255, 255), 2);
    }
  }

}  // namespace rm_auto_aim
